#include "BPatch_basicBlockLoop.h"
#include "BPatch_basicBlock.h"

#include "Symtab.h"
#include "Type.h"
#include "CFG.h"
#include "CodeSource.h"
#include "CodeObject.h"

#include "InstructionDecoder.h"
#include "Instruction.h"
#include "Operand.h"

#include "feature.h"
#include "graphlet.h"
#include "colors.h"
#include "supergraph.h"

#include "config.h"
#include "feature_analyzer.h"

#include <cstdio>
#include <cstring>

#include <string>
#include <vector>
#include <map>
#include <set>

#include <unordered_map>

#include <boost/iterator/filter_iterator.hpp>

using boost::make_filter_iterator;

using namespace std;
using namespace Dyninst;
using namespace Dyninst::ParseAPI;
using namespace Dyninst::SymtabAPI;
using namespace Dyninst::InstructionAPI;
using namespace graphlets;


static bool compound(Edge* e){
    static Intraproc intra_pred;
    static NoSinkPredicate nosink_pred;
    return intra_pred(e) && nosink_pred(e);
}

void FeatureAnalyzer::ProduceIdiomFeature(ParseAPI::Function* f, ParseAPI::Block *b){
    FeatureVector fv;
    map<string, int> counts;
	
    fv.eval(f,b, true,false);
	
    FeatureVector::iterator fvit = fv.begin();
    for( ; fvit != fv.end(); ++fvit) {
        counts[(*fvit)->format()] += 1;
    }
	
    map<string,int>::const_iterator cit = counts.begin();
    fprintf(config->idiomFile, "%lx" , b->start());
    for( ; cit != counts.end(); ++cit) {
        fprintf(config->idiomFile, " %s:%d",(*cit).first.c_str(),(*cit).second);
    }
    fprintf(config->idiomFile, "\n");
}

void FeatureAnalyzer::ProduceLoopIdiomFeature(const char *id, Dyninst::ParseAPI::Function * f, BPatch_basicBlockLoop *loop){
    map<string, pair<int, string> > counts;

    vector<BPatch_basicBlock*> bbs;
    loop->getLoopBasicBlocks(bbs);
    for (auto bit = bbs.begin(); bit != bbs.end(); ++bit) {
        ParseAPI::Block *b = ParseAPI::convert(*bit);
	FeatureVector fv;
	fv.eval(f,b, true,false);
	
	for (auto fvit = fv.begin(); fvit != fv.end(); ++fvit) {
	    const string &feature = (*fvit)->format();
	    if (counts.find(feature) == counts.end()) {
	        counts.insert(make_pair(feature, make_pair(1,  ((IdiomFeature*)(*fvit))->human_format())));
	    } else {
	        pair<int, string> &v = counts[feature];
		++v.first;
	    }
	}
    }
	
    fprintf(config->idiomFile, "%s" , id);
    for(auto cit = counts.begin() ; cit != counts.end(); ++cit) {
        AddAFeature(config->idiomFile, cit->first, cit->second.second, cit->second.first);
    }
    fprintf(config->idiomFile, "\n");
}




static graph * 
func_to_graph(BPatch_basicBlockLoop *loop, dyn_hash_map<Address,bool> & seen)
{
    dyn_hash_map<Address,snode*> node_map;

    int nctr = 0;

    vector<BPatch_basicBlock*> bbs;
    set<ParseAPI::Block*> blocks;
    loop->getLoopBasicBlocks(bbs);
    for (auto bit = bbs.begin(); bit != bbs.end(); ++bit) {
        ParseAPI::Block * b = ParseAPI::convert(*bit);
	blocks.insert(b);
    }
   

    graph * g = new graph();

    dyn_hash_map<size_t,bool> done_edges;
    
    for(auto bit = bbs.begin() ; bit != bbs.end(); ++bit) {
        Block * b = ParseAPI::convert(*bit);
        snode * n = g->addNode();

        char nm[16];
        snprintf(nm,16,"n%d",nctr++);
        n->name_ = std::string(nm);

        node_map[b->start()] = n;
        if(COLOR)
            n->setColor(new InsnColor(node_color(b)));
    }
    unsigned idx = 0;

    for(auto bit = bbs.begin() ; bit != bbs.end(); ++bit) {
        Block * b = ParseAPI::convert(*bit);

        if(seen.find(b->start()) != seen.end())
            continue;
        seen[b->start()] = true;

        snode * n = g->nodes()[idx];
        for(auto eit = make_filter_iterator(compound, b->sources().begin(), b->sources().end()); 
	         eit != make_filter_iterator(compound, b->sources().end(), b->sources().end()); 
		 ++eit) {
            Edge * e = *eit;
            if(done_edges.find((size_t)e) != done_edges.end())
                continue;
            done_edges[(size_t)e] = true;

            if(node_map.find(e->src()->start()) != node_map.end()) {
                (void)g->link(node_map[e->src()->start()],n,e->type());
            }
        }
        for(auto eit = make_filter_iterator(compound, b->targets().begin(), b->targets().end()); 
	         eit != make_filter_iterator(compound, b->targets().end(), b->targets().end()); 
		 ++eit) {
            Edge * e = *eit;
            if(done_edges.find((size_t)e) != done_edges.end())
                continue;
            done_edges[(size_t)e] = true;

            if(node_map.find(e->trg()->start()) != node_map.end()) {
                (void)g->link(n,node_map[e->trg()->start()],e->type());
            }
        }
        ++idx;
    }
    return g;
}

static void targets(Block * b, set<Block*> &t)
{   
    for(auto eit = make_filter_iterator(compound, b->sources().begin(), b->sources().end()); 
             eit != make_filter_iterator(compound, b->sources().end(), b->sources().end()); 
	     ++eit){
        t.insert((*eit)->trg());
    }
}

const int NULLDEF = numeric_limits<int>::max();
/*
 * Calls kill other call defs
 */
static void reaching_defs(
    ParseAPI::Function *f, 
    BPatch_basicBlockLoop *loop, 
    vector<Block*> & blocks, 
    vector< set<int> > & defs,
    dyn_hash_map<void*,int> & bmap)
{
    vector<int> work;
    vector<int> call_gen;
    Block * b = ParseAPI::convert(loop->getLoopHead());
    int bidx = bmap[b];

    // figure out the call generators
    call_gen.resize(blocks.size(),-1);
    for(auto cit = f->callEdges().begin() ; cit != f->callEdges().end(); ++cit) {
        if (bmap.find((*cit)->src()) == bmap.end()) continue;
        int cidx = bmap[(*cit)->src()];
        call_gen[ cidx ] = cidx;
    }

    // Even though the loop head may have a function call, we can still use NULLDEF to represent this callsite
    defs[bidx].insert(NULLDEF);
    work.push_back(bidx);

    while(!work.empty()) {
        bidx = work.back();
        work.pop_back();
        b = blocks[bidx];

        if(call_gen[bidx] != -1)
            defs[bidx].insert(call_gen[bidx]);

        set<Block*> targs;
        targets(b,targs);
        for(auto sit = targs.begin() ; sit != targs.end(); ++sit) {
            Block * t = *sit;
	    if (bmap.find(t) == bmap.end()) continue;
            int tidx = bmap[t];

            // if b makes a call, it kills all calls except that one
            // else it passes its defs
    
            set<int> newdefs = defs[tidx];
            if(call_gen[bidx] != -1)
                newdefs.insert(call_gen[bidx]);
            else
                newdefs.insert(defs[bidx].begin(),defs[bidx].end());

            if(newdefs != defs[tidx]) {
                defs[tidx] = newdefs;
                work.push_back(tidx);
            }
        }
    } 
}

/*
 * Construct a graph of entry, exit and call nodes, where edges indicate
 * reaching definitions
 */
static graph * collapse(
    ParseAPI::Function *f,
    BPatch_basicBlockLoop *loop, 
    vector<Block*> & blocks, 
    vector< set<int> > & defs,
    dyn_hash_map<void*,int> & bmap,
    dyn_hash_map<std::string,unsigned short> & libmap)
{
    vector<snode*> callnodes(blocks.size(),NULL);
    graph * g = new graph();
    snode * entry = g->addNode();

    // 1. Set up nodes for the call blocks
    for(auto cit = f->callEdges().begin() ; cit != f->callEdges().end(); ++cit) {
        if (bmap.find((*cit)->src()) == bmap.end()) continue;
        int cidx = bmap[(*cit)->src()];
        callnodes[ cidx ] = g->addNode();

        // color
        std::map<Address, std::string>::iterator pltit =
            f->obj()->cs()->linkage().find((*cit)->trg()->start());
        if(pltit != f->obj()->cs()->linkage().end()) {
            LibCallColor * c = new LibCallColor(libmap,(*pltit).second);
            callnodes[cidx]->setColor(c);
        } else {
            LocalCallColor * c = new LocalCallColor();
            callnodes[cidx]->setColor(c);
        }
    }

    // 2. Link the call nodes to one another
    for(unsigned i=0;i<callnodes.size();++i) {
        if(callnodes[i]) {
            set<int> & d = defs[i];
            for(auto it = d.begin() ; it != d.end(); ++it) {
                if(*it == NULLDEF)
                    g->link(entry,callnodes[i],0);
                else if(*it != (int)i)
                    g->link(callnodes[*it],callnodes[i],0);
            }
        }
    }

    // 3. Find exit nodes && link according to reaching defs       
    for(unsigned i=0;i<blocks.size();++i) {
        Block * b = blocks[i];
        int tcnt = 0;
	for(auto eit = make_filter_iterator(compound, b->targets().begin(), b->targets().end());                   
	         eit != make_filter_iterator(compound, b->targets().end(), b->targets().end()); 
	         ++eit){
             ++tcnt;
        }
        if(tcnt == 0) {
            snode * exit = g->addNode();
            set<int> & d = defs[i];
            for(auto it = d.begin() ; it != d.end(); ++it) {
                if(*it == NULLDEF)
                    g->link(entry,exit,0);
                else
                    g->link(callnodes[*it],exit,0);
            }
        }
    }
    return g;
}

static graph * mkcalldfa(ParseAPI::Function *f, BPatch_basicBlockLoop * loop,
    dyn_hash_map<std::string,unsigned short> & libmap)
{
    vector<Block*> blocks;
    vector< set<int> > defs; 
    dyn_hash_map<void*,int> bmap;

    vector<BPatch_basicBlock*> bbs;
    loop->getLoopBasicBlocks(bbs);

    for(auto bit = bbs.begin() ; bit != bbs.end(); ++bit) {
        ParseAPI::Block *b = ParseAPI::convert(*bit);
        bmap[b] = blocks.size();
        blocks.push_back(b);
    }
    defs.resize( blocks.size() );

    // 1. Reaching definitions on call blocks
    reaching_defs(f,loop, blocks,defs,bmap);
    
    // 2. Node collapse
   return collapse(f,loop, blocks,defs,bmap,libmap);
    
}


void FeatureAnalyzer::ProduceLoopGraphletFeature(const char* id, ParseAPI::Function* f, BPatch_basicBlockLoop *loop){
    map<graphlet,int> c1, c2, c3;
    dyn_hash_map<Address, bool> visited;
    graph * g = func_to_graph(loop,visited);
    g->mkgraphlets(c1, COLOR, ANON); // instruction graphlet

    // iteratively compress
    for(int m=0;m<MERGE;++m) {
        //unsigned sz = g->nodes().size();
        g->compact();
        //fprintf(stderr,"[%d] compacted graph from %d nodes to %ld nodes\n",
        //m,sz,g->nodes().size());
    }
        
    g->mkgraphlets(c2,COLOR,ANON); //supergraphlet
    delete g;

    dyn_hash_map<string, unsigned short> libmap;
    g = mkcalldfa(f, loop, libmap);
    g->mkgraphlets(c3, true, false);
    delete g;
    
    fprintf(config->graphletFile, "%s", id);
    for (auto cit=c1.begin(); cit != c1.end(); ++cit){
        AddAFeature(config->graphletFile, string("G_") + (cit->first).compact(COLOR), string("G_") + (cit->first).toString(), cit->second);
    }
    for (auto cit=c2.begin(); cit != c2.end(); ++cit){
        AddAFeature(config->graphletFile, string("S_") + (cit->first).compact(COLOR), string("S_") + (cit->first).toString(), cit->second);
    }
    for (auto cit=c3.begin(); cit != c3.end(); ++cit){
        AddAFeature(config->graphletFile, string("C_") + (cit->first).compact(COLOR), string("C_") + (cit->first).toString(), cit->second);
    }
    fprintf(config->graphletFile, "\n");
}

void FeatureAnalyzer::ProduceGraphletFeature(ParseAPI::Block* b){
    fprintf(config->graphletFile, "%lx", b->start());
}
void FeatureAnalyzer::ProduceLibcallFeature(CodeObject *co, ParseAPI::Function* f, int label){
    
    unordered_map<string, int> pltcnts;
    const ParseAPI::Function::edgelist & calls = f->callEdges();
    for(auto it = calls.begin() ; it != calls.end(); ++it) {
        Edge * e = *it;
        std::map<Address, std::string>::iterator pltit;
        pltit = co->cs()->linkage().find(e->trg()->start());
        if(pltit != co->cs()->linkage().end()) {
            pltcnts[(*pltit).second] += 1;
        }      
    }
    
    fprintf(config->libcallFile, "%d", label);
    for(auto cit = pltcnts.begin() ; cit != pltcnts.end(); ++cit) {
        if(internalFuncs.find((*cit).first) == internalFuncs.end()){
	    AddAFeature(config->libcallFile, cit->first, cit->first, cit->second);
        }
    }
    fprintf(config->libcallFile, "\n");

}

void FeatureAnalyzer::NGram(SymtabCodeSource *sts, ParseAPI::Block *b, unordered_map<string,int> &counts) {

    Address a;
    unsigned char ngrambuf[N];
    unsigned char * ngstart = ngrambuf;
    unsigned char * ngcur = ngrambuf;
    int cnt = 0;
        
    for(a = b->start(); a < b->end(); ++a) {
            
        unsigned char * byte = (unsigned char*)sts->getPtrToInstruction(a);
	*(ngcur++) = *byte;
	++cnt;
	
	if(ngcur - ngrambuf >= N)
	    ngcur = ngrambuf; 
	if(cnt >= N) {
	    char ngram[N+N+1];
	    unsigned char * tmp = ngstart++;
	    for (int i = 0; i < N; ++i){
	        sprintf(ngram+i+i, "%02x", *(tmp++));
		if (tmp - ngrambuf >= N)
		    tmp = ngrambuf;
	    }
            counts[string(ngram)] += 1;
            if(ngstart - ngrambuf >= N)
	        ngstart = ngrambuf; 
        }
    } 

}

void FeatureAnalyzer::ProduceNGramFeature(SymtabCodeSource* sts, ParseAPI::Block *b){
    unordered_map<string, int> counts;
   
    NGram(sts, b, counts);
    
    unordered_map<string, int>::iterator cit = counts.begin();
    fprintf(config->ngramFile, "%lx", b->start());
    for (; cit != counts.end(); ++cit){
        fprintf(config->ngramFile, " %s:%d", cit->first.c_str(), cit->second);
    }
    fprintf(config->ngramFile, "\n"); 
}

void FeatureAnalyzer::ProduceLoopNGramFeature(const char *id, SymtabCodeSource* sts, BPatch_basicBlockLoop *loop){
    unordered_map<string, int> counts;

    vector<BPatch_basicBlock*> bbs;
    loop->getLoopBasicBlocks(bbs);
    for (auto bit = bbs.begin(); bit != bbs.end(); ++bit) {
        ParseAPI::Block *b = ParseAPI::convert(*bit);
	NGram(sts, b, counts);
    } 
    
    fprintf(config->ngramFile, "%s", id);
    for (auto cit = counts.begin(); cit != counts.end(); ++cit){
        AddAFeature(config->ngramFile, cit->first, cit->first, cit->second);
    }
    fprintf(config->ngramFile, "\n"); 
}

void FeatureAnalyzer::ProduceDataFeature(ParseAPI::Function *f, Symtab * obj, int label){
    
    SymtabAPI::Function *func;
    if (!obj->findFuncByEntryOffset(func, f->addr())) {
        dprintf("cannot find SymtabAPI::Function for %s\n", f->name().c_str());
	return;
    }

    map<typeId_t, int> counts;
    map<typeId_t, int>::iterator cit;
    vector<localVar*> vars;
    vector<localVar*>::iterator vit;
    func->getLocalVariables(vars);

    for (vit = vars.begin(); vit != vars.end(); ++vit){
        counts[(*vit)->getType()->getID()] += 1;
    }
    fprintf(config->dataFile, "%d", label);
    for (cit = counts.begin(); cit != counts.end(); ++cit){
        fprintf(config->dataFile, " LOCAL_%d:%d", cit->first, cit->second);
    }

    counts.clear();
    vars.clear();

    func->getParams(vars);

    for (vit = vars.begin(); vit != vars.end(); ++vit){
        counts[(*vit)->getType()->getID()] += 1;     
    }
    for (cit = counts.begin(); cit != counts.end(); ++cit){
        fprintf(config->dataFile, " PARAM_%d:%d", cit->first, cit->second);
    }
    if (func->getReturnType() != NULL){
        fprintf(config->dataFile, " RET_%d:%d\n", func->getReturnType()->getID(), 1);
    } else {
        fprintf(config->dataFile, " RET_void:%d\n", 1);
    }

}

static int GetOperandScore(Operand& op) {
    
    set<RegisterAST::Ptr> regs;
    op.getReadSet(regs);
    op.getWriteSet(regs);

    if (op.readsMemory() || op.writesMemory()) {
	if (regs.size() == 0) 
	    return 3; // Memory access with address specified by an imm
	else
	    return 4; // Memory access with address specified by a reg
    }
    else {
	if (regs.size() == 0) 
	    return 1; // Access an imm 
	else
	    return 2; // Access a reg
    }
    return 0;
}

int FeatureAnalyzer::AddressingMode(Dyninst::ParseAPI::Block *b){ 
    int score = 0;

    Instruction::Ptr insn;
    Address cur = b->start();
    void * buf  = b->region()->getPtrToInstruction(b->start());
    if(!buf) {
        return 0;
    }
    
    InstructionDecoder dec((unsigned char*)buf,b->size(), b->region()->getArch());
    while((insn = dec.decode())) {
        vector<Operand> operands;
	insn->getOperands(operands);
	for (auto opIt = operands.begin(); opIt != operands.end(); ++opIt)
	    score += GetOperandScore(*opIt);

        cur += insn->size();
    }										    
    return score;
}

void FeatureAnalyzer::ProduceAddressingModeFeature(Dyninst::ParseAPI::Block *b){
    int score = AddressingMode(b);
    fprintf(config->addressingModeFile, "%lx ADDRESSING:%d\n", b->start(), score);
}
 
void FeatureAnalyzer::ProduceLoopAddressingModeFeature(const char* id, BPatch_basicBlockLoop *loop) {
    int score = 0;
    vector<BPatch_basicBlock*> bbs;
    loop->getLoopBasicBlocks(bbs);
    for (auto bit = bbs.begin(); bit != bbs.end(); ++bit) {
        ParseAPI::Block *b = ParseAPI::convert(*bit);
	score += AddressingMode(b);
    } 
    fprintf(config->addressingModeFile, "%s", id);
    AddAFeature(config->addressingModeFile, "ADDRESSING", "ADDRESSING", score);
    fprintf(config->addressingModeFile, "\n");
}

void FeatureAnalyzer::AddAFeature(FILE* f, const string &feature, const string &human_format, int value) {
    if (featureIndex.find(feature) == featureIndex.end()) {
        ++totalFeatures;
	featureIndex.insert(make_pair(feature, make_pair(totalFeatures, human_format)));
    }
    fprintf(f, " %d:%d", featureIndex[feature].first, value);
}

